We might want to consider adding this to requirements-dev.txt.

Yes, unless there is a compelling interest to keeping this separate, in which case a requirements-jax.txt or something could be a middle ground. That would centralize the prereqs (with pinned versions if necessary) so it does not have to be explained in multiple places but can still be skipped in, e.g., environments that must run tests but would have a problem installing JAX.

The readthedocs check passed and has always passed installing only the dependencies in https://github.com/pymc-devs/pymc/blob/main/conda-envs/environment-dev.yml (should be the same as in requirements-dev.txt). I think the error you mention in thr issue is different. One reason might be differemces between the conda and txt requirements, maybe jax is missing from the txt? jax is needed to import anything in sampling jax file, but numpyro and blackjax are both optional

One reason might be differemces between the conda and txt requirements, maybe jax is missing from the txt? jax is needed to import anything in sampling jax file, but numpyro and blackjax are both optional

So jax is definitely missing from requirements-dev.txt, but I assumed that was done on purpose because it is specifically excluded in scripts/generate_pip_deps_from_conda.py (introduced via #5477).

Edit to add, the squashed commit has some possible rationale:

* Add jax to dev environment

* Jax doesn't work on windows, remove it

* Exclude jax

I think we should continue to exclude jax dependencies from requirements-dev.txt until they are available on all platforms.

But why are you just adding numpyro to the instructions and not blackjax?

Update, @isms: in #5938 I propose to remove the docs build script for Windows, and also added a note in the paragraph above the code block you're changing.

I'd just like to understand if we should also include blackjax here?

@michaelosthege Installing blackjax is not necessary to make the documentation build once numpyro is installed, that's the only reason not to include both. Happy to add it back if that feels important.

Thanks @isms !